// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/pollux_exception.h>

#include <melon/synchronization/atomic_struct.h>
#include <exception>

namespace kumo {
    namespace pollux {
        std::exception_ptr toPolluxException(const std::exception_ptr &exceptionPtr) {
            try {
                std::rethrow_exception(exceptionPtr);
            } catch (const PolluxException &) {
                return exceptionPtr;
            } catch (const std::exception &e) {
                return std::make_exception_ptr(
                    PolluxUserError(std::current_exception(), e.what(), false));
            }
        }

        int64_t &threadNumPolluxThrow() {
            thread_local int64_t numThrow = 0;
            return numThrow;
        }

        bool &threadSkipErrorDetails() {
            thread_local bool skipErrorDetails{false};
            return skipErrorDetails;
        }

        ExceptionContext &getExceptionContext() {
            thread_local ExceptionContext context;
            return context;
        }

        // Traverses the context hierarchy and appends messages from all contexts that
        // are marked as essential.
        std::string getAdditionalExceptionContextString(
            PolluxException::Type exceptionType,
            const std::string &currentMessage) {
            auto *context = &getExceptionContext();
            std::string additionalMessage = "";
            if (!context->parent || !context->parent->parent) {
                return additionalMessage;
            }
            context = context->parent;
            while (context->parent) {
                if (context->isEssential) {
                    auto message = context->message(exceptionType);
                    if (!message.empty()) {
                        additionalMessage += message + " ";
                    }
                }
                context = context->parent;
            }
            if (!additionalMessage.empty()) {
                // Get rid of the extra space at the end.
                additionalMessage.pop_back();
            }
            return additionalMessage;
        }

        PolluxException::PolluxException(
            const char *file,
            size_t line,
            const char *function,
            std::string_view failingExpression,
            std::string_view message,
            std::string_view errorSource,
            std::string_view errorCode,
            bool isRetriable,
            Type exceptionType,
            std::string_view exceptionName)
            : PolluxException(State::make(exceptionType, [&](auto &state) {
                state.exceptionType = exceptionType;
                state.exceptionName = exceptionName;
                state.file = file;
                state.line = line;
                state.function = function;
                state.failingExpression = failingExpression;
                state.message = message;
                state.errorSource = errorSource;
                state.errorCode = errorCode;
                state.context = getExceptionContext().message(exceptionType);
                state.additionalContext =
                        getAdditionalExceptionContextString(exceptionType, state.context);
                state.isRetriable = isRetriable;
            })) {
        }

        PolluxException::PolluxException(
            const std::exception_ptr &e,
            std::string_view message,
            std::string_view errorSource,
            std::string_view errorCode,
            bool isRetriable,
            Type exceptionType,
            std::string_view exceptionName)
            : PolluxException(State::make([&](auto &state) {
                state.exceptionType = exceptionType;
                state.exceptionName = exceptionName;
                state.file = "UNKNOWN";
                state.line = 0;
                state.function = "";
                state.failingExpression = "";
                state.message = message;
                state.errorSource = errorSource;
                state.errorCode = errorCode;
                state.context = getExceptionContext().message(exceptionType);
                state.additionalContext =
                        getAdditionalExceptionContextString(exceptionType, state.context);
                state.isRetriable = isRetriable;
                state.wrappedException = e;
            })) {
        }

        namespace {
            /// returns whether PolluxException stacktraces are enabled and whether, if they
            /// are rate-limited, whether the rate-limit check passes
            bool isStackTraceEnabled(PolluxException::Type type) {
                using namespace std::literals::chrono_literals;
                const bool isSysException = type == PolluxException::Type::kSystem;
                if ((isSysException && !turbo::get_flag(FLAGS_pollux_exception_system_stacktrace_enabled)) ||
                    (!isSysException && !turbo::get_flag(FLAGS_pollux_exception_user_stacktrace_enabled))) {
                    // PolluxException stacktraces are disabled.
                    return false;
                }

                const int32_t rateLimitMs = isSysException
                                                ? turbo::get_flag(
                                                    FLAGS_pollux_exception_system_stacktrace_rate_limit_ms)
                                                : turbo::get_flag(FLAGS_pollux_exception_user_stacktrace_rate_limit_ms);
                // not static so the gflag can be manipulated at runtime
                if (0 == rateLimitMs) {
                    // PolluxException stacktraces are not rate-limited
                    return true;
                }
                static melon::AtomicStruct<std::chrono::steady_clock::time_point> systemLast;
                static melon::AtomicStruct<std::chrono::steady_clock::time_point> userLast;
                auto *last = isSysException ? &systemLast : &userLast;

                auto const now = std::chrono::steady_clock::now();
                auto latest = last->load(std::memory_order_relaxed);
                if (now < latest + std::chrono::milliseconds(rateLimitMs)) {
                    // PolluxException stacktraces are rate-limited and the rate-limit check
                    // failed
                    return false;
                }

                // PolluxException stacktraces are rate-limited and the rate-limit check
                // passed
                //
                // the cas happens only here, so the rate-limit check in effect gates not
                // only computation of the stacktrace but also contention on this atomic
                // variable
                return last->compare_exchange_strong(latest, now, std::memory_order_relaxed);
            }
        } // namespace

        template<typename F>
        std::shared_ptr<const PolluxException::State> PolluxException::State::make(
            PolluxException::Type exceptionType,
            F f) {
            auto state = std::make_shared<PolluxException::State>();
            if (isStackTraceEnabled(exceptionType)) {
                // new v.s. make_unique to avoid any extra frames from make_unique
                state->stackTrace.reset(new process::StackTrace());
            }
            f(*state);
            return state;
        }

        /*
        Not much to say. Constructs the elaborate message from the available
        pieces of information.
         */
        void PolluxException::State::finalize() const {
            assert(elaborateMessage.empty());

            // Fill elaborateMessage_
            if (!exceptionName.empty()) {
                elaborateMessage += "Exception: ";
                elaborateMessage += exceptionName;
                elaborateMessage += '\n';
            }

            if (!errorSource.empty()) {
                elaborateMessage += "Error Source: ";
                elaborateMessage += errorSource;
                elaborateMessage += '\n';
            }

            if (!errorCode.empty()) {
                elaborateMessage += "Error Code: ";
                elaborateMessage += errorCode;
                elaborateMessage += '\n';
            }

            if (!message.empty()) {
                elaborateMessage += "Reason: ";
                elaborateMessage += message;
                elaborateMessage += '\n';
            }

            elaborateMessage += "Retriable: ";
            elaborateMessage += isRetriable ? "True" : "False";
            elaborateMessage += '\n';

            if (!failingExpression.empty()) {
                elaborateMessage += "Expression: ";
                elaborateMessage += failingExpression;
                elaborateMessage += '\n';
            }

            if (!context.empty()) {
                elaborateMessage += "Context: " + context + "\n";
            }

            if (!additionalContext.empty()) {
                elaborateMessage += "Additional Context: " + additionalContext + "\n";
            }

            if (function) {
                elaborateMessage += "Function: ";
                elaborateMessage += function;
                elaborateMessage += '\n';
            }

            if (file) {
                elaborateMessage += "File: ";
                elaborateMessage += file;
                elaborateMessage += '\n';
            }

            if (line) {
                elaborateMessage += "Line: ";
                auto len = elaborateMessage.size();
                size_t t = line;
                do {
                    elaborateMessage += static_cast<char>('0' + t % 10);
                    t /= 10;
                } while (t);
                reverse(elaborateMessage.begin() + len, elaborateMessage.end());
                elaborateMessage += '\n';
            }

            elaborateMessage += "Stack trace:\n";
            if (stackTrace) {
                elaborateMessage += stackTrace->toString();
            } else {
                elaborateMessage += "Stack trace has been disabled.";
                if (exceptionType == PolluxException::Type::kSystem) {
                    elaborateMessage +=
                            " Use --pollux_exception_system_stacktrace_enabled=true to enable it.\n";
                } else {
                    elaborateMessage +=
                            " Use --pollux_exception_user_stacktrace_enabled=true to enable it.\n";
                }
            }
        }

        const char *PolluxException::State::what() const noexcept {
            try {
                melon::call_once(once, [&] { finalize(); });
                return elaborateMessage.c_str();
            } catch (...) {
                return "<unknown failure in PolluxException::what>";
            }
        }
    } // namespace pollux
} // namespace kumo
